package cn.jly.flink.state;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

/**
 * 状态编程
 *
 * @author lanyangji
 */
public class CountWindowAverageStateApp {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.fromElements(
                Tuple2.of("tom", 4),
                Tuple2.of("tom", 2),
                Tuple2.of("tom", 5),
                Tuple2.of("tom", 6),
                Tuple2.of("tom", 4)
        )
                .keyBy(0)
                .flatMap(new CountWindowAverage("prefix"))
                .print();

        env.execute();
    }

    static class CountWindowAverage extends RichFlatMapFunction<Tuple2<String, Integer>, Tuple2<String, Integer>> {

        private String prefix;

        /**
         * valueState的用法，第一个是count, 第二个是sum
         */
        private transient ValueState<Tuple2<Integer, Integer>> countAndSum;

        public CountWindowAverage(String prefix) {
            this.prefix = prefix;
        }

        public String getPrefix() {
            return prefix;
        }

        public void setPrefix(String prefix) {
            this.prefix = prefix;
        }

        @Override
        public void flatMap(Tuple2<String, Integer> value, Collector<Tuple2<String, Integer>> out) throws Exception {
            // 访问valueState的值
            Tuple2<Integer, Integer> countAndSumVal = countAndSum.value();

            // update count
            countAndSumVal.f0 += 1;

            // update sum
            countAndSumVal.f1 += value.f1;

            // 更新状态
            countAndSum.update(countAndSumVal);

            // 如果 count 等于 2, 写出平均值并清除状态
            if (countAndSumVal.f0 >= 2) {
                out.collect(Tuple2.of(this.prefix + "-" + countAndSumVal.f0, countAndSumVal.f1 / countAndSumVal.f0));
                countAndSum.clear();
            }
        }

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            ValueStateDescriptor<Tuple2<Integer, Integer>> descriptor
                    = new ValueStateDescriptor<>("countAndSum", TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>() {
            }), Tuple2.of(0, 0));
            countAndSum = getRuntimeContext().getState(descriptor);
        }
    }
}

